/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2020-2021. All rights reserved.
 * Description: 梯度调优模式集合通信算子图优化头文件
 * Author: lilianlin
 * Create: 2020-09-10
 */

#ifndef AUTO_TUNING_HCOM_GRAPH_OPTIMIZER_H
#define AUTO_TUNING_HCOM_GRAPH_OPTIMIZER_H

#include "../hcom/hcom_graph_optimizer.h"
#include "common/optimizer/graph_optimizer.h"
#include "common/optimizer/graph_optimizer_types.h"
#include "graph/compute_graph.h"
#include "base.h"

namespace hccl {

class AutoTuningHcomGraphOptimizer : public HcomGraphOptimizer {
public:
    AutoTuningHcomGraphOptimizer();
    ~AutoTuningHcomGraphOptimizer() override;
    ge::Status Initialize(const map<std::string, std::string> &options,
                          ge::OptimizeUtility *const optimizeUtility) override;
    // optimize original graph, using in graph preparation stage
    ge::Status OptimizeOriginalGraph(ge::ComputeGraph &graph) override;
    ge::Status OptimizeFusedGraph(ge::ComputeGraph& graph) override;
protected:
    HcclResult CheckSupportedOP(const std::string &sCollectiveType) const override;
    HcclResult CalcOpRunningParam(ge::Node& node);
    HcclResult SetOpOutputMemSize(ge::Node& node, const std::string &sCollectiveType) override;
    HcclResult CalcHCCLOutputMemSize(const std::string &sCollectiveType, int64_t& memSize) override;
    HcclResult SetOpMemAttr(ge::Node &node, const std::string &sCollectiveType, const u64 &opMemSize) override;

private:
    std::string workPath_;
    bool isGradientAutoTune_;
};
}

#endif